import math
from QKDNetwork import QKDNetwork
from docplex.mp.model import Model

class LP:
    def __init__(self, net: QKDNetwork):
        self.net = net
        self.node_num = net.num_nodes + 1 # 用户方便使用ID索引节点
        self.sd_num = len(net.sd_list)
        self.max_candidate_size = net.max_candidate_size
        self.selected_sole_paths = {}
        self.select_paths()
        
    def select_paths(self):
        self.processed_selected_paths_index = []
        for i, _ in enumerate(self.net.sd_list):
            self.process_path(i)
        # print("处理完成:", self.processed_selected_paths_index)
            
    def process_path(self, index):
        model = Model(name="rps_he")
        M = 1000
        
        q_dim2 = self.max_candidate_size
        q = model.binary_var_dict(
            [j for j in range(q_dim2)],
            name="q"
        )
        
        p_dim1, p_dim2 = self.node_num, self.node_num
        p = model.integer_var_dict(
            [(i, j) for i in range(p_dim1) for j in range(p_dim2)],
            lb=0,
            name="p"
        )
        T = model.integer_var(lb=0, name="T")
        
        for u in range(self.net.num_nodes):
            for v in self.net.G.neighbors(u):
                model.add_constraint(p[u, v] <= M * self.net.G.nodes[u]["transmitter"])
        
        for u in range(self.net.num_nodes):
            p_sum = model.linear_expr()
            for v in self.net.G.neighbors(u):
                p_sum += p[u, v]
            model.add_constraint(p_sum <= T)
        
        q_sum = model.linear_expr()
        for j in range(len(self.net.candidate_paths[index]["paths"])):
            q_sum += q[j]
        model.add_constraint(q_sum == 1)
        
        def is_sublist(sub, main):
            sub_length = len(sub)
            return any(sub == main[i:i+sub_length] for i in range(len(main) - sub_length + 1))
        
        for edge in self.net.G.edges:
            u = edge[0]
            v = edge[1]
            flow = model.linear_expr()
            flow += self.net.G.nodes[u]["transmitter_rate"] * p[u, v]
            flow += self.net.G.nodes[v]["transmitter_rate"] * p[v, u]
            for j, path in enumerate(self.net.candidate_paths[index]["paths"]):
                if(is_sublist([u, v], path) or is_sublist([v, u], path)):
                    # (u, v)在链路上
                    flow -= self.net.sd_list[index][2] * q[j]
            for sd_index, sd_selected_path_index in enumerate(self.processed_selected_paths_index):
                path = self.net.candidate_paths[sd_index]["paths"][sd_selected_path_index]
                if(is_sublist([u, v], path) or is_sublist([v, u], path)):
                    flow -= self.net.sd_list[sd_index][2]
            model.add_constraint(flow >= 0)
        
        model.minimize(T)
        sol = model.solve()
        if sol:
            for i, variable in q.items():
                if(variable.solution_value == 1):
                    self.processed_selected_paths_index.append(i)
                    # print(f"append:{i}")
        
        
                
    